% =========================================================================
% Test code for Natural Image Matting using Deep Convolutional Neural Networks (CNN matting)
%
% Donghyeon Cho
% EE Department, KAIST
% cdh12242@gmail.com
% =========================================================================

clear all;
%% settings
run('./function/vlfeat-0.9.16/toolbox/vl_setup.m');
addpath(genpath('./function/matting'));
folder_img = './data/inputs';
model = './test1.mat';
flag = true;
for fol=1:3    
    savepath_Closed = ['./result/ClosedForm_matting' num2str(fol)];
    savepath_KNN = ['./result/KNN_matting' num2str(fol)];
    savepath_CNN = ['./result/CNN_matting' num2str(fol)];
    folder_trimap = ['./data/Trimap' num2str(fol)];
    filepaths_img = dir(fullfile(folder_img,'*.png'));
    filepaths_trimap = dir(fullfile(folder_trimap,'*.png'));

    if ~exist(savepath_Closed,'dir')
        mkdir(savepath_Closed);
    end
    
    if ~exist(savepath_KNN,'dir')
        mkdir(savepath_KNN);
    end
    
    if ~exist(savepath_CNN,'dir')
        mkdir(savepath_CNN);
    end

    %% read image
    for i=1:length(filepaths_img)
        disp(['Trimap  = ' num2str(fol) ', Image  = ' filepaths_img(i).name(1:end-4)]);
        image = im2double(imread(fullfile(folder_img,filepaths_img(i).name)));
        trimap = im2double(imread(fullfile(folder_trimap,filepaths_trimap(i).name)));
        if size(trimap,3)==3
            trimap = rgb2gray(trimap);
        end
        trimap(trimap>0.99)=1;trimap(trimap<0.01)=0;trimap(trimap<0.8&trimap>0.2)=0.5;

        %% closed form matting
        consts_map = trimap;
        consts_map(:) = 0;
        consts_map(trimap==1|trimap==0) = 1;
        consts_vals = consts_map;
        consts_vals(trimap==0) =0;
        matte_closed=solveAlphaC2F(image,consts_map,consts_vals,1, ...
                                    1,[],[],[]);
        %% KNN matting                        
        matte_knn = knnmatting(image,trimap); 
        %% Normalized RGB        
        image = image./repmat(sqrt(image(:,:,1).^2+image(:,:,2).^2+image(:,:,3).^2),[1 1 3]);
        %% setting input to CNN
        input(:,:,1:3) = image-0.5;
        input(:,:,4) = matte_closed-0.5;
        input(:,:,5) = matte_knn-0.5;
        %% forward CNN 
        tic
        [matte_CNN] = Matting_forward(model, input);
        toc
        
        matte_CNN = matte_CNN +0.5;    
        matte_CNN(trimap==1) = 1;
        matte_CNN(trimap==0) = 0;
        matte_CNN(matte_CNN>1) = 1;
        matte_CNN(matte_CNN<0) = 0;
        % if flag is true, pixels which have different estimated alpha values from the Closed form and KNN matting will be only considered
        if flag
            matte_CNN(abs(matte_closed-matte_knn)<0.05) = matte_closed(abs(matte_closed-matte_knn)<0.05);
        end
        %% save results
        imwrite(matte_CNN,[savepath_CNN '/' filepaths_img(i).name]);  
        imwrite(matte_closed,[savepath_Closed '/' filepaths_img(i).name]);
        imwrite(matte_knn,[savepath_KNN '/' filepaths_img(i).name]);
        
        clear input; 
    end
   
end